package SecureChat;
import java.io.*;
import java.net.*;
import java.math.*;
import java.util.*;
import SecureChat.Convert;

public class DHAuth {
        
    public static byte[] Client(BigInteger g, String IDa, String IDb, Socket s)
    {
        //declare variables
        byte[] key = {0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00};
        byte[] receivedHash = {0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00};
        byte[] expectedHash = {0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00,
			       0x00,0x00,0x00,0x00};
        byte [] readBuffer;
        BigInteger a,p,ga,gb,gab;
        String response = null;
        InputStream sIn;
        OutputStream sOut;
        
        //Check debug mode
        boolean DEBUG;
        if(System.getProperty("DEBUG") != null){
            DEBUG = true;
        }
        else{
            DEBUG = false;
        }

        
        if(DEBUG){
            System.err.println("Client: Starting authentication in client mode");
        }
        
        //make call to rand. num gen to get a;
        a = Utilities.GenRandNum();

        if(DEBUG){
            System.err.println("DEBUG: Random number generated");
        }

        // make call to rand prime gen to get p;
        p = Utilities.GenRandPrime();

        if(DEBUG){
            System.err.println("DEBUG: Random prime generated"); 
        }

        //calculate g^a mod p
        ga = Utilities.CalcModPower(g, a, p);

        if(DEBUG){
            System.err.println("DEBUG: G^a mod p calculated");
        }

        //Create input streams
        try{
            sIn = s.getInputStream();
        }
        catch(IOException e){
            System.out.print("Could not open input stream: " + e);
            return null;
        }

        //Create output streams
        try{
            sOut = s.getOutputStream();
        }
        catch(IOException e){
            System.out.print("Could not open output stream: " + e);
            return null;
        }
                        
        if(DEBUG){
            System.err.println("DEBUG: Beginning authentication protocol.");
            System.err.println("DEBUG: Sending prime");
        }
        
        try{
            sOut.write(p.toByteArray());
        }
        catch(IOException e){
            System.out.print("Failed to write prime to output stream: " + e);
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Waiting for response");
        }

        //listen for OK response
        while(true){
            try{ 
                if (sIn.available() > 0){       
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    if(DEBUG){
                        System.err.println("DEBUG: Received response.");
                    }
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed: " + e);
                return null;
            }
        }

        if(readBuffer[0] != 1 || readBuffer.length > 1){
            System.out.println("Server protocol failure: Fail response received");
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Sending G^a mod p");
        }
        
        try{
            sOut.write(ga.toByteArray());
        }

        catch(IOException e){
            System.out.print("Could not write g^a to output stream: " + e);
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Waiting for server response");
        }
        //now, Server should successfully have p, and ga
        
        //so we now listen for the server's response - receivedHash
        while(true){
            try{ 
                if (sIn.available() > 0){
                    if(DEBUG){
                        System.err.println("DEBUG: Response received (receivedHash)");
                    }
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed: " + e);
                return null;
            }
        }
        
        //get the hash
        receivedHash = readBuffer;
        
        //let server know successfully received hash
        try{
            sOut.write(1);
        }
        catch(IOException e){
            System.out.print("Failed to write to output stream: " + e);
            return null;
        }

        //so we now listen for the server's response - bobs gb
        while(true){
            try{ 
                if (sIn.available() > 0){
                    if(DEBUG){
                        System.err.println("DEBUG: Response received (g^b mod p)");
                    }
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);                         
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed on input stream: " +e);
                return null;
            }
        }
        
        gb = new BigInteger(readBuffer);
        
        //now we want to calculate g^ab mod p
        gab = Utilities.CalcModPower(gb, a, p);
        if(DEBUG){
            System.err.println("DEBUG: Verifying hash from server with calculated hash");
        }
        
        //now we verify with the H[g^ab mod p] that we received from bob...
        try{
            expectedHash = Utilities.HashValues(gab,IDb);
        }
        catch(Exception e){
            System.out.println("Utilities.HashValues() failed: " + e);
            return null;
        }

	if(DEBUG){
            System.err.println("DEBUG: Expected hash = " + Convert.toHex(expectedHash));
            System.err.println("DEBUG: Received hash = " + Convert.toHex(receivedHash));
        }
        if(Arrays.equals(expectedHash, receivedHash)){
            if(DEBUG){
                System.err.println("DEBUG: Hashes are equivalent.");
            }
                
            //since verified, we send Hash[g^abmodp, "alice"]
            try{
                receivedHash = Utilities.HashValues(gab, IDa);
            }
            catch(Exception e){
                System.out.println("Utilities.HashValues() failed: " + e);
                return null;
            }
            try{
                sOut.write(receivedHash);
            }
            catch(IOException e){
                System.out.print("Failed to write to output stream: " + e);
                return null;
            }
            key = gab.toByteArray();
        }
        else{
            System.out.println("Authentication failed. Agreed key inconsistent.");
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Authentication successful; Key = " + Convert.toHex(key));
        }

        return key;
    }
    
    public static byte[] Server(BigInteger g, String IDa, String IDb, Socket s)
    {
        byte[] key = {0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00,
                      0x00,0x00,0x00,0x00};
        byte[] hash = {0x00,0x00,0x00,0x00,
                       0x00,0x00,0x00,0x00,
                       0x00,0x00,0x00,0x00,
                       0x00,0x00,0x00,0x00};
	byte [] readBuffer;
	byte [] receivedHash;
	byte [] expectedHash;
		
        BigInteger p,b,ga,gb,gab;
        String response = null;
        String ID = null;
        String ok = "ok";

        //Check debug mode
        boolean DEBUG;
        if(System.getProperty("DEBUG") != null){
            DEBUG = true;
        }
        else{
            DEBUG = false;
        }
        
        if(DEBUG){
            System.err.println("DEBUG: Starting authentication in server mode");
        }
        
        //make call to rand. num gen to get b;
        b = Utilities.GenRandNum();

        if(DEBUG){
            System.err.println("DEBUG: Random number generated");
        }       
        
        InputStream sIn;
        OutputStream sOut;
        
        //Create socket connection
        try{
            sIn = s.getInputStream();
        }
        catch(IOException e){
            System.out.print("Could not open input stream: " + e);
            return null;
        }
           
        try{
            sOut = s.getOutputStream();
        }
        catch(IOException e){
            System.out.print("Could not open output stream: " + e);
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Beginning authentication protocol.");
            System.err.println("DEBUG: Waiting for prime from client.");
        }
        
        //get the prime p
        while(true){
            try{ 
                if (sIn.available() > 0){
                    if(DEBUG){
                        System.err.println("DEBUG: Received response.");
                    }
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed on input stream: " + e);
                return null;
            }
        }

        p = new BigInteger(readBuffer);

        if(DEBUG){
            System.err.println("DEBUG: Sending prime acknowledgement and waiting for reply.");
        }        

        try{
            sOut.write(1);
        }
        catch(IOException e){
            System.out.print("Write failed on output stream: " + e);
            return null;
        }

        //get ga
        while(true){
            try{ 
                if (sIn.available() > 0){
                    if(DEBUG){
                        System.err.println("DEBUG: g^a mod p received.");
                    }    
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    //response = new String(readBuffer);
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed" +e);
                return null;
            }
        }
        
        
        ga = new BigInteger(readBuffer);

        if(DEBUG){
            System.err.println("DEBUG: Calculating g^ab mod p.");
        }        

        //now we have IDA, ga, and p.
        //we will now calculate g^ab mod p and send it
        gab = Utilities.CalcModPower(ga, b, p);
        // hash it
        try{
            hash = Utilities.HashValues(gab, IDb);
        }
        catch(Exception e){
            System.out.println("Failure in Utilities.HashValue(): " + e);
            return null;
        }
        if(DEBUG){
            System.err.println("DEBUG: Sending hash = " + Convert.toHex(hash));
        }        
        try{
            sOut.write(hash);
        }
        catch(IOException e){
            System.out.print("Failed to write to output stream: " + e);
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Waiting for acknowledgement response");
        }

        //listen for OK response
        while(true){
            try{ 
                if (sIn.available() > 0){
                        
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    //response = new String(readBuffer);
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed" +e);
                return null;
            }
        }
        if(readBuffer[0] != 1 || readBuffer.length > 1){
            System.out.println("Client protocol failure: Fail response receieved");
            return null;
        }
        
        //calculate gb and send it
        if(DEBUG){
            System.err.println("DEBUG: Calculating g^b mod p and sending to client");
        }
        gb = Utilities.CalcModPower(g,b,p);
        
        try{
            sOut.write(gb.toByteArray());
        }
        
        catch(IOException e){
            System.out.print("Failed to write to output stream" + e);
            return null;
        }
                
        //alice will verify, listen here for her gab value
        while(true){
            try{ 
                if (sIn.available() > 0){
                    if(DEBUG){
                        System.err.println("DEBUG: Response received");
                    }
                    readBuffer = new byte[sIn.available()];
                    sIn.read(readBuffer);
                    //response = new String(readBuffer);
                    break;
                }
            } 
            catch (IOException e){
                System.out.println("Read failed" + e);
                return null;
            }
        }
        receivedHash = readBuffer;
        
        try{
            expectedHash = Utilities.HashValues(gab, IDa);
        }
        catch(Exception e){
            System.out.println("Failure in Utilites.HashValues(): " + e);
            return null;
        }
        if(DEBUG){
            System.err.println("DEBUG: Expected hash = " + Convert.toHex(expectedHash));
            System.err.println("DEBUG: Received hash = " + Convert.toHex(receivedHash));
        }
        
        if(Arrays.equals(expectedHash, receivedHash)){
            key = gab.toByteArray();
        }
        else{
            System.out.println("Authentication failed. Agreed key inconsistent.");
            return null;
        }

        if(DEBUG){
            System.err.println("DEBUG: Authentication successful; Key = " + Convert.toHex(key));
        }
        return key; 
        //key will either be empty (same as initialized byte array) 
        //or the gab value.
    }
}//end class

